﻿using System;
using System.Collections.Generic;
using System.Text;

namespace MLForgeSharp.Models.ProbabilisticModels.Bayes
{
    /// <summary>
    /// 贝叶斯信念网
    /// </summary>
    public class BayesNet
    {
        public List<BayesNode> Nodes { get; set; }

        public BayesNet()
        {
            Nodes = new List<BayesNode>();
        }

        public void AddNode(BayesNode node)
        {
            Nodes.Add(node);
        }

        public BayesNode GetNodeByName(string name)
        {
            return Nodes.Find(n => n.Name == name);
        }

        public double CalculateJointProbability(Dictionary<string, bool> evidence)
        {
            double jointProbability = 1.0;

            foreach (var node in Nodes)
            {
                string condition = "";
                foreach (var parent in node.Parents)
                {
                    condition += evidence[parent.Name] ? "T" : "F";
                }

                double probability = node.GetConditionalProbability(condition);
                jointProbability *= evidence[node.Name] ? probability : 1 - probability;
            }

            return jointProbability;
        }
    }

    /// <summary>
    /// 贝叶斯信念网节点
    /// </summary>
    public class BayesNode
    {
        public string Name { get; set; }
        public List<BayesNode> Parents { get; set; }
        public List<BayesNode> Children { get; set; }
        public Dictionary<string, double> ConditionalProbabilities { get; set; }

        public BayesNode(string name)
        {
            Name = name;
            Parents = new List<BayesNode>();
            Children = new List<BayesNode>();
            ConditionalProbabilities = new Dictionary<string, double>();
        }

        public void AddParent(BayesNode parent)
        {
            Parents.Add(parent);
            parent.Children.Add(this);
        }

        public void SetConditionalProbability(string condition, double probability)
        {
            ConditionalProbabilities[condition] = probability;
        }

        public double GetConditionalProbability(string condition)
        {
            return ConditionalProbabilities.ContainsKey(condition) ? ConditionalProbabilities[condition] : 0.0;
        }
    }

    public class BayesNetExample
    {
        public BayesNetExample()
        {
            // 创建节点
            BayesNode rain = new BayesNode("Rain");
            BayesNode sprinkler = new BayesNode("Sprinkler");
            BayesNode grassWet = new BayesNode("GrassWet");

            // 设置条件概率表
            rain.SetConditionalProbability("", 0.2); // P(Rain = T)
            sprinkler.SetConditionalProbability("", 0.1); // P(Sprinkler = T)
            grassWet.SetConditionalProbability("TT", 0.99); // P(GrassWet = T | Rain = T, Sprinkler = T)
            grassWet.SetConditionalProbability("TF", 0.9);  // P(GrassWet = T | Rain = T, Sprinkler = F)
            grassWet.SetConditionalProbability("FT", 0.8);  // P(GrassWet = T | Rain = F, Sprinkler = T)
            grassWet.SetConditionalProbability("FF", 0.0);  // P(GrassWet = T | Rain = F, Sprinkler = F)

            // 构建网络
            grassWet.AddParent(rain);
            grassWet.AddParent(sprinkler);

            // 创建贝叶斯网络
            BayesNet network = new BayesNet();
            network.AddNode(rain);
            network.AddNode(sprinkler);
            network.AddNode(grassWet);

            // 设置证据
            Dictionary<string, bool> evidence = new Dictionary<string, bool>
            {
                { "Rain", true },
                { "Sprinkler", false },
                { "GrassWet", true }
            };

            // 计算联合概率
            double jointProbability = network.CalculateJointProbability(evidence);
            Console.WriteLine("Joint Probability: " + jointProbability);
        }
    }
}
